package cn.claim.storage.client;

import cn.claim.storage.tools.Config;

import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

/**
 * Created by claim on 17-7-7.
 */

/**
 * File crypter for AES algorithm.
 * The tool class will create a .tmp file to a temporary directory.
 */
public final class Crypter {
    private static Map<String, Cipher> ciphers = new HashMap<>();

    private static Cipher getCipherInstance(String id, int cipherMode) {
        ObjectInputStream reader = null;
        if (ciphers.get(id) != null) return ciphers.get(id + cipherMode);
        else {
            Cipher cipher = null;
            try {
                byte[] buffer = new byte[256];
                File keyFile = new File("./key_" + id);
                if (!keyFile.exists())
                    flushLocalAESKey(id);
                reader = new ObjectInputStream(new FileInputStream(keyFile));
                int length = reader.read(buffer);
                SecretKeySpec secretKey = new SecretKeySpec(Arrays.copyOf(buffer, length), "AES");
                cipher = Cipher.getInstance("AES");
                cipher.init(cipherMode, secretKey);
                ciphers.put(id + cipherMode, cipher);
            } catch (NoSuchAlgorithmException | InvalidKeyException | NoSuchPaddingException | IOException e) {
                e.printStackTrace();
            } finally {
                try {
                    if (reader != null) {
                        reader.close();
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            return cipher;
        }
    }

    /**
     * May use {@code ./key_${id}} as the AES key file or create a new key
     * (if the key file not exists).
     *
     * @param id A identifier for the key you used.
     * @return A File object referred to the encrypted file with a suffix of ".encrypted".
     */
    public static File encryptFile(File sourceFile, String id) {
        File encryptedFile = null;
        InputStream inputStream = null;
        OutputStream outputStream = null;
        try {
            inputStream = new FileInputStream(sourceFile);
            encryptedFile = File.createTempFile("___", ".encrypted");
            outputStream = new FileOutputStream(encryptedFile);
            Cipher cipher = getCipherInstance(id, Cipher.ENCRYPT_MODE);
            CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
            byte[] buffer = new byte[Integer.parseInt(Config.getProp("bufferSize", "1024"))];
            int nRead = 0;
            while ((nRead = cipherInputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, nRead);
                outputStream.flush();
            }
            cipherInputStream.close();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (inputStream != null) {
                    inputStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (outputStream != null) {
                    outputStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return zipFile(encryptedFile);
    }

    /**
     * May use {@code ./key_${id}} as the AES key file or create a new key
     * (if the key file not exists).
     *
     * @param id A identifier for the key you used.
     * @return A File object referred to the decrypted file
     */
    public static File decryptFile(File sourceFile, String id, File target) {
        sourceFile = unzipFile(sourceFile);
        InputStream inputStream = null;
        OutputStream outputStream = null;
        try {
            Cipher cipher = getCipherInstance(id, Cipher.DECRYPT_MODE);
            inputStream = new FileInputStream(sourceFile);
            outputStream = new FileOutputStream(target);
            CipherOutputStream cipherOutputStream = new CipherOutputStream(outputStream, cipher);
            byte[] buffer = new byte[Integer.parseInt(Config.getProp("bufferSize", "1024"))];
            int r;
            while ((r = inputStream.read(buffer)) >= 0) {
                cipherOutputStream.write(buffer, 0, r);
            }
            cipherOutputStream.flush();
            cipherOutputStream.close();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (inputStream != null) {
                    inputStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (outputStream != null) {
                    outputStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return target;
    }

    public static File zipFile(File source) {
        File result = null;
        InputStream inputStream = null;
        ZipOutputStream zip = null;
        try {
            result = File.createTempFile("___", ".zip");
            inputStream = new FileInputStream(source);
            byte[] buffer = new byte[Integer.parseInt(Config.getProp("bufferSize", "1024"))];
            zip = new ZipOutputStream(new FileOutputStream(result));
            zip.putNextEntry(new ZipEntry("/zippedFile"));
            int length;
            while ((length = inputStream.read(buffer)) != -1) {
                zip.write(buffer, 0, length);
            }
            zip.flush();

        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (inputStream != null) {
                    inputStream.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (zip != null) {
                    zip.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return result;
    }

    public static File unzipFile(File source) {
        File result = null;
        FileInputStream input = null;
        FileOutputStream file = null;
        try {
            input = new FileInputStream(source);
            ZipInputStream zip = new ZipInputStream(input);
            zip.getNextEntry();
            result = File.createTempFile("___", ".unzip");
            file = new FileOutputStream(result);
            byte[] buffer = new byte[Integer.parseInt(Config.getProp("bufferSize", "1024"))];
            int length;
            while ((length = zip.read(buffer)) != -1) {
                file.write(buffer, 0, length);
            }
            file.flush();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (input != null) {
                    input.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (file != null) {
                    file.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return result;
    }

    /**
     * The method will flush the key you saved in {@code ./config}.
     * <em>Pay attention! If you didn't save your key manually before invoke this method,
     * you may be never able to decrypt the file you encrypted.</em>
     */
    public static void flushLocalAESKey(String id) {
        KeyGenerator keyGenerator;
        ObjectOutputStream writer = null;
        try {
            keyGenerator = KeyGenerator.getInstance("AES");
            keyGenerator.init(128, new SecureRandom());
            SecretKey secretKey = keyGenerator.generateKey();
            File keyFile = new File("./key_" + id);
            if (!keyFile.exists())
                //noinspection ResultOfMethodCallIgnored
                keyFile.createNewFile();
            writer = new ObjectOutputStream(new FileOutputStream(keyFile));
            writer.write(secretKey.getEncoded());
        } catch (NoSuchAlgorithmException | IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (writer != null) {
                    writer.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * The method should be used to encode a file encrypted by {@link #encryptFile(File, String)}.
     *
     * @return A {@code File} Object refers to a temp file with the suffix of ".base64".
     */
    public static File toBase64(File sourceFile) {
        InputStream reader = null;
        PrintWriter writer = null;
        Base64.Encoder base64Encoder = Base64.getEncoder();
        byte[] buffer = new byte[3 * Integer.parseInt(Config.getProp("bufferSize", "1024"))];
        File tempFile = null;
        try {
            tempFile = File.createTempFile("___", ".base64");
            reader = new FileInputStream(sourceFile);
            writer = new PrintWriter(tempFile);
            int length;
            while ((length = reader.read(buffer)) != -1) {
                writer.print(base64Encoder.encodeToString(Arrays.copyOf(buffer, length)));
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (reader != null)
                    reader.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
            if (writer != null) {
                writer.close();
            }
        }
        return tempFile;
    }

    /**
     * Should be used to generate a file can be decrypted by{@link #decryptFile(File, String, File)}
     *
     * @param base64File A file encoded by {@link #toBase64(File)}.
     * @return A temp file encrypted by {@link #encryptFile(File, String)}.
     */
    public static File fromBase64(File base64File) {
        InputStream reader = null;
        OutputStream writer = null;
        File tempFile = null;
        Base64.Decoder base64Decoder = Base64.getDecoder();
        try {
            int length;
            tempFile = File.createTempFile("___", ".encrypted");
            reader = new FileInputStream(base64File);
            writer = new FileOutputStream(tempFile);
            byte[] buffer = new byte[Integer.parseInt(Config.getProp("bufferSize", "1024"))];
            while ((length = reader.read(buffer)) != -1) {
                byte[] temp = base64Decoder.decode(Arrays.copyOf(buffer, length));
                writer.write(temp);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (reader != null) {
                    reader.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
            try {
                if (writer != null)
                    writer.close();
            } catch (IOException e) {
                e.printStackTrace();
            }

        }
        return tempFile;
    }

//    /**
//     * Only for test
//     */
//    public static void main(String[] args) {
//        File file = encryptFile(new File("./data/core"), "default");
//        try {
//            System.out.println(file.getCanonicalPath());
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//        file = toBase64(file);
//        try {
//            System.out.println(file.getCanonicalPath());
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//        file = fromBase64(file);
//        try {
//            System.out.println(file.getCanonicalPath());
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//        decryptFile(file,
//                "default", new File("./data/a"));
//    }
}
